#include <opencv2/opencv.hpp>
#include "utils/image_processing/image_process.h"
#include "txda_runtime.h"

int main(int argc, char **argv) {
  if (argc < 3) {
    std::cerr << "Usage: " << argv[0] << " <image_path> <model_path>" << std::endl;
  }

  std::string image_path = argv[1];

  // Load the image
  cv::Mat src = cv::imread(image_path);

  // 预处理
  std::vector<int64_t> input_shape = {1, 3, 640, 640}; // batch, channel, height, width
  auto input_tensor_values = tcv::PreprocessImage<std::uint16_t>(src, input_shape);
  std::cout << "[DEBUG] Input tensor values size: " << input_tensor_values.size() << std::endl;

  std::string model_path = argv[2];
  const char* logfile = nullptr;

  // 1. create txdart instance
  auto txdart = txda::TxdaRT(logfile, txda::ILogger::Level::Info);

  // 2. get logger ptr
  auto logger_ = txdart.txdartGetLogger();

  // 3. init runtime instance
  logger_->log(txda::ILogger::Level::Debug, "** Start init runtime instance **");
  txdart.txdartInitRuntime();

  // 4. set user using devices
  uint32_t tp_size = 1;
  txda::TxdaResult result = txda::TxdaResult::TXDA_SUCCESS;

  uint32_t device_num = 1;
  logger_->log(txda::ILogger::Level::Debug, "** Start set device idx **");
  std::vector<uint32_t> card_setting = txdart.txdartGetEnvCardSetting(device_num);
  if (card_setting.empty()) {
    logger_->log(txda::ILogger::Level::Error, "TXDA_VISIBLE_DEVICES not set!");
  }
  txdart.txdartSetDeviceIdx(result, card_setting, tp_size);

  // 5. init txda context
  std::string option = "-O2";
  bool enable_compile = false;
  logger_->log(txda::ILogger::Level::Debug, "** Start init context **");
  txdart.txdartCtxinit(result, tp_size, model_path, option, enable_compile);
  if (txda::TxdaResult::TXDA_SUCCESS != result) {
    logger_->log(txda::ILogger::Level::Error, "Context init failed!");
    txdart.txdartDeinitRuntime();
    return -1;
  }

  // 6. compile and load model
  logger_->log(txda::ILogger::Level::Debug, "** Start compileAndLoadModel **");
  txdart.txdartCompileAndLoadModel(result);
  if (txda::TxdaResult::TXDA_SUCCESS != result) {
    logger_->log(txda::ILogger::Level::Error, "Compile and load model failed!");
    return -1;
  }

  /* read param */
  logger_->log(txda::ILogger::Level::Debug, "** Start read weight file to host mem **");

  auto &txda_chip_info = txdart.txdartGetChipInfo(0, 0);
  auto param_size = txda_chip_info->getParamSize(0);
  auto param_ptr = txda::utils::read_file(model_path + "/chip0/agent/data/param.bin", param_size);
  txda_chip_info->addParamHostAddr((uint64_t)param_ptr);

  std::vector<uint8_t*> in_ptrs{(uint8_t*)input_tensor_values.data()};
  if (in_ptrs.size() != txda_chip_info->getInputNum()) {
    logger_->log(txda::ILogger::Level::Error, "Inputs num not equal to model info's definition.");
    return -1;
  } else {
    for (uint32_t i = 0; i < txda_chip_info->getInputNum(); i++) {
      txda_chip_info->addInputHostAddr((uint64_t)in_ptrs[i]);
    }
  }

  std::vector<uint16_t> infer_output;
  infer_output.resize(1*600*3);
  std::vector<uint8_t*> out_ptrs{(uint8_t*)infer_output.data()};
  if (out_ptrs.size() != txda_chip_info->getOutputNum()) {
    logger_->log(txda::ILogger::Level::Error, "Outputs num not equal to model info's definition.");
    return -1;
  } else {
    for (uint32_t i = 0; i < txda_chip_info->getOutputNum(); i++) {
      txda_chip_info->addOutputHostAddr((uint64_t)out_ptrs[i]);
    }
  }

  // 7. malloc weight device mem
  logger_->log(txda::ILogger::Level::Debug, "** Start malloc mem for weight **");
  if (txda::TxdaResult::TXDA_SUCCESS != txdart.txdartMemMalloc(true)) {
    logger_->log(txda::ILogger::Level::Error, "Malloc weight mem failed!");
    return -1;
  }

  // 8. copy weight to device mem 
  logger_->log(txda::ILogger::Level::Debug, "** Start copy weight to device **");
  if (txda::TxdaResult::TXDA_SUCCESS != txdart.txdartMemcpyHtoD(true)) {
    logger_->log(txda::ILogger::Level::Error, "Copy weight to device failed!");
    return -1;
  }

  txda_chip_info->freeParamHostAddr();

  // 9. create boot param
  logger_->log(txda::ILogger::Level::Debug, "** Start create boot param **");
  if (txda::TxdaResult::TXDA_SUCCESS != txdart.txdartCreateBootParam()) {
    logger_->log(txda::ILogger::Level::Error, "Create boot param failed!");
    return -1;
  }

  // 10. malloc in/output device mem
  logger_->log(txda::ILogger::Level::Debug, "** Start malloc mem for input & output **");
  if (txda::TxdaResult::TXDA_SUCCESS != txdart.txdartMemMalloc(false)) {
    logger_->log(txda::ILogger::Level::Error, "Malloc input & output mem failed!");
    return -1;
  }

  // 11. set boot param
  logger_->log(txda::ILogger::Level::Debug, "** Start set boot param **");
  if (txda::TxdaResult::TXDA_SUCCESS != txdart.txdartSetBootParam()) {
    logger_->log(txda::ILogger::Level::Error, "Txad set boot param failed!");
    return -1;
  }

  logger_->log(txda::ILogger::Level::Debug, "** Start copy input to device **");
  if (txda::TxdaResult::TXDA_SUCCESS != txdart.txdartMemcpyHtoD(false)) {
    return -1;
  }

  logger_->log(txda::ILogger::Level::Debug, "** Start launch kernel **");
  if (txda::TxdaResult::TXDA_SUCCESS != txdart.txdartLaunchKernel()) {
    logger_->log(txda::ILogger::Level::Error, "Txad launch kernel failed!");
    return -1;
  }

  logger_->log(txda::ILogger::Level::Debug, "** Start copy output to host **");
  if (txda::TxdaResult::TXDA_SUCCESS != txdart.txdartMemcpyDtoH()) {
    logger_->log(txda::ILogger::Level::Error, "Txad copy output to host failed!");
    return -1;
  }

  cv::Mat out_infer_mat(1, static_cast<int>(infer_output.size()), CV_16FC1);
  memcpy(out_infer_mat.data, infer_output.data(), infer_output.size() * sizeof(uint16_t));
  cv::Mat out_mat;
  out_infer_mat.convertTo(out_mat, CV_32FC1);
  std::vector<float> out_vec;
  out_vec.assign(out_mat.ptr<float>(0), out_mat.ptr<float>(0) + out_mat.total());

  float confidence_threshold = 0.6;
  std::vector<tcv::Detection> detections = tcv::FilterDetections(out_vec, 
    confidence_threshold, input_shape[3], input_shape[2], src.cols, src.rows);

  cv::Mat result_image = tcv::DrawLabels(src, detections);

  // Save or display the result image
  cv::imwrite("output_inference.jpg", result_image);

  // Context destroy
  txdart.txdartCtxDestroy();

  return 0;
}